import requests
from openai import OpenAI

from AI接口.base import APIClient


class DeepSeekClient(APIClient):
    def __init__(self, api_key):
        super().__init__(base_url="https://api.deepseek.com", api_key=api_key)
        self.client = OpenAI(api_key=api_key, base_url=self.base_url)

    def chat(self, messages, model="deepseek-chat", stream=False, **kwargs):
        return self.client.chat.completions.create(
            model=model,
            messages=messages,
            stream=stream,
            **kwargs
        )

    def get_result(self, response):
        """从response中获取结果"""
        return response.choices[0].message.content


if __name__ == '__main__':
    from settings import DEEPSEEK_API_KEY

    deepseek_client = DeepSeekClient(api_key=DEEPSEEK_API_KEY)
    # response = deepseek_client.chat(
    #     messages=[
    #         {"role": "user", "content": "你是什么模型？"},
    #     ]
    # )
    # print(deepseek_client.get_result(response))

    # JSON格式调用
    system_prompt = """
    The user will provide some exam text. Please parse the "novels" and output them in JSON format. 

    EXAMPLE INPUT: 
    生成几个和腕表有关的小说和小说简介.

    EXAMPLE JSON OUTPUT:
    {
        "novels": [
            {
                "title": "时间囚笼：我的手表能重启人生",
                "genre": "都市奇幻",
                "description": "林默从未想过，一块破旧的机械表会彻底改变他的人生。那是一个阴雨绵绵的午后，他在古董店避雨时，老板神秘兮兮地将这块表塞给了他。表盘上刻着古怪的纹路，指针永远停在午夜十二点。
第一次转动表冠是在他目睹女友车祸的瞬间。时间倒流，他成功救下了她，却发现母亲因此突发心脏病离世。第二次，他挽救了母亲，却导致整个街区陷入火灾。每一次重启，都像是打开了一个潘多拉魔盒，命运的齿轮在看不见的地方悄然转动。
更可怕的是，他发现自己的记忆开始出现裂痕，有时分不清现实与重启前的世界。而那块表，似乎正在吸取他的生命力，表盘上的纹路越来越像他掌心的纹路。
在无数次的时间循环中，林默逐渐意识到，这不是命运的馈赠，而是一场精心设计的囚笼。表盘上的十二个刻度，暗示着他仅剩的十二次重启机会。而最后一次重启，等待他的将是永恒的困局。"
            },
        ]
    }
    """
    response = deepseek_client.chat(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": "生成一个和腕表有关的小说和小说简介,简介在250-300自之间，title参考例子的结构，但是不要和例子一样。"},
        ]
    )
    print(deepseek_client.get_result(response))
